import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# 定义一个简单的循环神经网络（RNN）模型用于文本生成
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        super(SimpleRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden):
        out, hidden = self.rnn(x, hidden)
        out = self.fc(out[:, -1, :])
        return out, hidden

    def init_hidden(self, batch_size):
        return torch.zeros(self.num_layers, batch_size, self.hidden_size)

# 模型参数
input_size = 10  # 输入特征维度
hidden_size = 20 # 隐藏层大小
output_size = 5  # 输出词汇表大小
num_layers = 1   # RNN层数
learning_rate = 0.001
num_epochs = 10  # 训练轮数
batch_size = 64  # 批量大小
sequence_length = 10  # 序列长度

# 实例化模型、损失函数和优化器
model = SimpleRNN(input_size, hidden_size, output_size, num_layers)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 假设inputs和targets是准备好的数据，inputs的形状为(batch_size, sequence_length, input_size)
# targets的形状为(batch_size, sequence_length)

# 训练循环
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for i in range(0, inputs.size(0) - 1, sequence_length):
        inputs_batch = inputs[i:i+sequence_length]
        targets_batch = targets[i:i+sequence_length]

        # 初始化隐藏状态
        hidden = model.init_hidden(batch_size)

        optimizer.zero_grad()

        # 前向传播
        outputs, hidden = model(inputs_batch, hidden)

        # 设置温度参数（例如，temperature=1.0表示不调整，小于1.0使输出更尖锐，大于1.0使输出更平坦）
        temperature = 1.2  # 示例温度值
        logits = outputs / temperature

        # 计算损失
        loss = criterion(logits.view(-1, output_size), targets_batch.view(-1))
        loss.backward()

        # 更新参数
        optimizer.step()

        total_loss += loss.item()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(inputs)}')

